import platform

import torch
from torch.utils.data import Dataset, DataLoader
import torchvision
from torchvision import transforms, datasets
from setting import batch_size, cpu_workers


# 定义数据集
class MyDataset(Dataset):

    def __init__(self, train):
        # 在线加载数据集
        self.data = datasets.CIFAR100(root='../data', train=train, download=True)

        # 数据增强
        self.compose = transforms.Compose([
            transforms.Resize(300),  # 改变尺寸
            transforms.RandomHorizontalFlip(p=0.5),  # 随机左右翻转
            transforms.ToTensor(),  # 图像转矩阵数据，0-1
            # 让三个通道的数据分别服从3个正态分布
            transforms.Normalize(mean=[0.845, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def __len__(self):
        return len(self.data) // 1000 if platform.system() == 'Windows' else len(self.data)

    def __getitem__(self, i):
        x, y = self.data[i]
        x = self.compose(x)
        return x, y


# 回调函数，每次从loader取一批数据时回调，可以在这里做一些数据整理工作
# def collect_fn(data):


# print("开始加载训练数据集")
train_dataset = MyDataset(train=True)
train_dataloader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, drop_last=True,
                              num_workers=cpu_workers,
                              # collate_fn=collate_fn
                              )
# print("训练数据加载完成", len(train_dataloader))


# print("开始加载测试数据集")
test_dataset = MyDataset(train=False)
test_dataloader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False, drop_last=False,
                             num_workers=cpu_workers,
                             # collate_fn=collate_fn
                             )
# print("测试数据加载完成", len(test_dataloader))


# if __name__ == '__main__':
#     x, y = next(iter(train_dataloader))
#     print(x.shape, y.shape)
#     x, y = next(iter(test_dataloader))
#     print(x.shape, y.shape)
